"""
Spatial maps of total expained variance 
by HYCOM in SWOT SCIENCE orbit

B. Yadidya
Feb 25, 2026
"""

import sys
sys.path.append('.')
from common_imports import *

dir = "data/"
ssh_ds = xr.open_dataset(dir + 'regridded_0.25deg_ALL_SSH_VR_v2.0.1.nc')
sss_ds = xr.open_dataset(dir + 'regridded_0.25deg_ALL_SLOPE_VR_v2.0.1.nc')

# Decrease resolution from 4km to 0.25 deg for plotting
sss_x = (sss_ds.fsss_var_x_hycom_no_eddies)
sss_y = (sss_ds.fsss_var_y_hycom_no_eddies)
ssh = (ssh_ds.fssh_var_hycom_no_eddies)

"""
# I used the 4km regridded datasets for the plots shown in the manuscript. 
# However for storage regions, 0.25 deg regridded datasets are uploaded on Dataverse.
ssh_ds = xr.open_dataset(dir + 'regridded_4km_ALL_SSH_VR_v2.0.1.nc')
sss_ds = xr.open_dataset(dir + 'regridded_4km_ALL_SLOPE_VR_v2.0.1.nc')

# Decrease resolution from 4km to 0.25 deg for plotting
sss_x = coarsen_data(sss_ds.fsss_var_x_hycom_no_eddies, .1)
sss_y = coarsen_data(sss_ds.fsss_var_y_hycom_no_eddies, .1)
ssh = coarsen_data(ssh_ds.fssh_var_hycom_no_eddies, .1)
"""
fig_width_cm = 18.4
fig_width_in = fig_width_cm / 2.54
fig_height_max_cm = 19.0
fig_height_in = fig_height_max_cm / 2.54 

fig, axs = plt.subplots(
    3, 1, 
    figsize=(fig_width_in, fig_height_in), 
    constrained_layout=True
)
axs = axs.ravel()

# --- 4. Plotting ---
limit1 = 4
levels1 = np.linspace(-limit1, limit1, 21)
plot1 = ssh.plot.contourf(ax=axs[0], levels=levels1, 
                         x='longitude', cmap=colormaps.prinsenvlag_r,
                         cbar_kwargs={'shrink': .7, 'pad': .01})

# Update colorbar to 8pt and ensure line weights
plot1.colorbar.set_label('cm$^2$', size=8)
plot1.colorbar.set_ticks([-limit1, 0, limit1])
plot1.colorbar.ax.tick_params(which='major', labelsize=8, size=0, width=0) 
plot1.colorbar.ax.minorticks_off()
plot1.colorbar.outline.set_linewidth(0.5)

limit2 = 8
plot2 = (sss_x/1e-9).plot.contourf(ax=axs[1], levels=np.linspace(-limit2, limit2, 21),
                               x='longitude', cmap=colormaps.prinsenvlag_r,
                               cbar_kwargs={'shrink': .7, 'pad': .01})

plot3 = (sss_y/1e-9).plot.contourf(ax=axs[2], levels=np.linspace(-limit2, limit2, 21),
                               x='longitude', cmap=colormaps.prinsenvlag_r,
                               cbar_kwargs={'shrink': .7, 'pad': .01})     

# Update colorbar 2
plot2.colorbar.set_ticks([-limit2, 0, limit2])
plot2.colorbar.ax.tick_params(which='major', labelsize=8, size=0, width=0)
plot2.colorbar.ax.minorticks_off()
plot2.colorbar.set_label('10$^{-3}$(cm/km)$^2$', size=8)
plot2.colorbar.outline.set_linewidth(0.5)

# Update colorbar 3
plot3.colorbar.set_label('10$^{-3}$(cm/km)$^2$', size=8)
plot3.colorbar.set_ticks([-limit2, 0, limit2])
plot3.colorbar.ax.tick_params(which='major', labelsize=8, size=0, width=0)
plot3.colorbar.ax.minorticks_off()
plot3.colorbar.outline.set_linewidth(0.5)


# --- 5. Final Formatting ---
axs[0].set_title('Total explained SSH variance by HYCOM', fontsize=9)
axs[1].set_title('Total explained cross-track slope variance by HYCOM', fontsize=9)
axs[2].set_title('Total explained along-track slope variance by HYCOM', fontsize=9)    

for ax in axs:
    plot_global_map_on_axes_basemap(ax) 
    ax.set_aspect('equal')
    ax.set_xlabel('')
    ax.set_ylabel('')
    ax.set_ylim(-60, 60)

# --- 6. Panel Labels with Specific Bold Font ---
panel_label_props = dict(
    fontproperties=bold_font_props,  # <--- Uses the loaded Bold OTF file
    fontsize=9,
    ha='center',
    va='center',
    bbox=dict(
        boxstyle='round, pad=0.1',
        facecolor='white',
        edgecolor='none',
        alpha=1
    )
)

for i, label in enumerate(['A', 'B', 'C']):
    axs[i].text(0.017, .94, label, transform=axs[i].transAxes, **panel_label_props)

plt.savefig('Fig7_HYCOM_total_explained_variance.png', dpi=500, bbox_inches='tight')
